{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "291deab3",
   "metadata": {},
   "source": [
    "# 来源：https://zhaozhiming.github.io/2024/06/01/llamaindex-llama3-es-hybrid-search/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bb4a4e1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.base.embeddings.base import BaseEmbedding, Embedding\n",
    "from llama_index.core.bridge.pydantic import PrivateAttr\n",
    "from typing import Any, List"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "dcdf6140",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "def send_request(model_uid: str, text: str, url: str):\n",
    "    url = f\"{url}/v1/embeddings\"\n",
    "    request_body = {\"model\": model_uid, \"input\": text}\n",
    "    response = requests.post(url, json=request_body)\n",
    "    if response.status_code != 200:\n",
    "        raise RuntimeError(\n",
    "            f\"Failed to create the embeddings, detail: {response.text}\"\n",
    "        )\n",
    "    return response.json()\n",
    "\n",
    "def get_embedding(text: str, model_uid: str, url: str) -> Embedding:\n",
    "    \"\"\"Get embedding.\"\"\"\n",
    "    text = text.replace(\"\\n\", \" \")\n",
    "    response_data = send_request(model_uid, text, url)\n",
    "    return response_data[\"data\"][0][\"embedding\"]\n",
    "\n",
    "def get_embeddings(\n",
    "    list_of_text: List[str], model_uid: str, url: str\n",
    ") -> List[Embedding]:\n",
    "    \"\"\"Get embeddings.\"\"\"\n",
    "    assert len(list_of_text) <= 2048, \"The batch size should not be larger than 2048.\"\n",
    "\n",
    "    list_of_text = [text.replace(\"\\n\", \" \") for text in list_of_text]\n",
    "    response_data = send_request(model_uid, list_of_text, url)\n",
    "    return [d[\"embedding\"] for d in response_data[\"data\"]]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c2f984f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomEmbeddings(BaseEmbedding):\n",
    "    \"\"\"Custom class for embeddings.\n",
    "\n",
    "    Args:\n",
    "        model_name (str): Mode for embedding.\n",
    "        url(str): Url for embedding model.\n",
    "    \"\"\"\n",
    "\n",
    "    _model_name: str = PrivateAttr()\n",
    "    _url: str = PrivateAttr()\n",
    "\n",
    "    def __init__(self, model_name: str, url: str, **kwargs: Any) -> None:\n",
    "        self._model_name = model_name\n",
    "        self._url = url\n",
    "        super().__init__(**kwargs)\n",
    "\n",
    "    @classmethod\n",
    "    def class_name(cls) -> str:\n",
    "        return \"custom_embedding\"\n",
    "\n",
    "    def _aget_query_embedding(self, query: str) -> Embedding:\n",
    "        return get_embedding(text=query, model_uid=self._model_name, url=self._url)\n",
    "\n",
    "    def _aget_text_embedding(self, text: str) -> Embedding:\n",
    "        return get_embedding(text=text, model_uid=self._model_name, url=self._url)\n",
    "\n",
    "    def _get_query_embedding(self, query: str) -> Embedding:\n",
    "        return get_embedding(text=query, model_uid=self._model_name, url=self._url)\n",
    "\n",
    "    def _get_text_embedding(self, text: str) -> Embedding:\n",
    "        return get_embedding(text=text, model_uid=self._model_name, url=self._url)\n",
    "\n",
    "    def _get_text_embeddings(self, texts: List[str]) -> List[Embedding]:\n",
    "        return get_embeddings(\n",
    "            list_of_text=texts, model_uid=self._model_name, url=self._url\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "716ea613",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 数据入库\n",
    "\n",
    "from llama_index.vector_stores.elasticsearch import ElasticsearchStore\n",
    "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext\n",
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "\n",
    "store = ElasticsearchStore(\n",
    "    index_name=\"a\",\n",
    "    es_url=\"http://localhost:9200\",\n",
    ")\n",
    "documents = SimpleDirectoryReader(\"./data\").load_data()\n",
    "node_parser = SentenceSplitter(chunk_size=256, chunk_overlap=50)\n",
    "storage_context = StorageContext.from_defaults(vector_store=store)\n",
    "embed_model = CustomEmbeddings(\n",
    "    model=\"BAAI/bge-base-en-v1.5\", url=\"http://localhost:6006\"\n",
    ")\n",
    "VectorStoreIndex.from_documents(\n",
    "    documents,\n",
    "    transformations=[node_parser],\n",
    "    embed_model=embed_model,\n",
    "    storage_context=storage_context,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a901cff7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 全文检索\n",
    "from llama_index.vector_stores.elasticsearch import AsyncBM25Strategy\n",
    "from llama_index.core import Settings\n",
    "\n",
    "text_store = ElasticsearchStore(\n",
    "    index_name=\"a\",\n",
    "    es_url=\"http://localhost:9200\",\n",
    "    retrieval_strategy=AsyncBM25Strategy(),\n",
    ")\n",
    "Settings.embed_model = embed_model\n",
    "text_index = VectorStoreIndex.from_vector_store(\n",
    "    vector_store=text_store,\n",
    ")\n",
    "text_retriever = text_index.as_retriever(similarity_top_k=2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e4ae0bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 向量检索\n",
    "from llama_index.vector_stores.elasticsearch import AsyncDenseVectorStrategy, AsyncSparseVectorStrategy\n",
    "\n",
    "vector_store = ElasticsearchStore(\n",
    "    index_name=\"a\",\n",
    "    es_url=\"http://localhost:9200\",\n",
    "    retrieval_strategy=AsyncDenseVectorStrategy(),\n",
    "    # retrieval_strategy=AsyncSparseVectorStrategy(model_id=\".elser_model_2\"),\n",
    ")\n",
    "Settings.embed_model = embed_model\n",
    "vector_index = VectorStoreIndex.from_vector_store(\n",
    "    vector_store=vector_store,\n",
    ")\n",
    "vector_retriever = vector_index.as_retriever(similarity_top_k=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "025a611d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 混合检索\n",
    "from llama_index.vector_stores.elasticsearch import AsyncDenseVectorStrategy\n",
    "\n",
    "vector_store = ElasticsearchStore(\n",
    "    index_name=\"avengers\",\n",
    "    es_url=\"http://localhost:9200\",\n",
    "    retrieval_strategy=AsyncDenseVectorStrategy(hybrid=True),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed5aa8b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 混合检索\n",
    "from typing import List\n",
    "from llama_index.core.schema import NodeWithScore\n",
    "\n",
    "def fuse_results(results_dict, similarity_top_k: int = 2):\n",
    "    \"\"\"Fuse results.\"\"\"\n",
    "    k = 60.0\n",
    "    fused_scores = {}\n",
    "    text_to_node = {}\n",
    "\n",
    "    # 计算倒数排名分数\n",
    "    for nodes_with_scores in results_dict.values():\n",
    "        for rank, node_with_score in enumerate(\n",
    "            sorted(\n",
    "                nodes_with_scores, key=lambda x: x.score or 0.0, reverse=True\n",
    "            )\n",
    "        ):\n",
    "            text = node_with_score.node.get_content()\n",
    "            text_to_node[text] = node_with_score\n",
    "            if text not in fused_scores:\n",
    "                fused_scores[text] = 0.0\n",
    "            fused_scores[text] += 1.0 / (rank + k)\n",
    "\n",
    "    # 结果按分数排序\n",
    "    reranked_results = dict(\n",
    "        sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)\n",
    "    )\n",
    "\n",
    "    # 结果还原为节点集合\n",
    "    reranked_nodes: List[NodeWithScore] = []\n",
    "    for text, score in reranked_results.items():\n",
    "        reranked_nodes.append(text_to_node[text])\n",
    "        reranked_nodes[-1].score = score\n",
    "\n",
    "    return reranked_nodes[:similarity_top_k]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
